"""
基于Agno框架和AkShare数据接口的AI投资分析系统

工作流主程序
"""

import os
import sys
import argparse
import logging
import atexit
import signal
from datetime import datetime, timedelta
from typing import Optional
from dotenv import load_dotenv

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from src.agents import MarketDataAgent, TechnicalAnalystAgent, FundamentalAnalystAgent
from src.utils.cotools import _save_collected_data_to_file
from agno.workflow.v2 import Workflow, Step, Steps,Parallel
from agno.workflow.v2.types import StepOutput, StepInput

# 加载环境变量
load_dotenv()

# 确保日志目录存在（在配置日志之前）
os.makedirs("logs", exist_ok=True)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('logs/investment_agent.log'),
        logging.StreamHandler()
    ]
)

logger = logging.getLogger(__name__)


def cleanup_on_exit():
    """程序退出时的清理函数"""
    try:
        # 强制刷新所有输出缓冲区
        sys.stdout.flush()
        sys.stderr.flush()

        # 给Rich库一些时间完成清理
        import time
        time.sleep(0.1)
    except:
        pass  # 忽略清理过程中的任何错误


def signal_handler(signum, frame):
    """信号处理函数"""
    logger.info("收到退出信号，正在清理...")
    cleanup_on_exit()
    sys.exit(0)


# 注册退出处理函数
atexit.register(cleanup_on_exit)
signal.signal(signal.SIGINT, signal_handler)
signal.signal(signal.SIGTERM, signal_handler)

class InvestmentAnalysisSystem:
    """投资分析系统主类"""

    def __init__(self, model_id: str = "deepseek-chat"):
        """
        初始化投资分析系统，并提前初始化智能体实例
        """
        self.model_id = model_id
        # 提前初始化智能体（Agent 对象）
        self.market_data_agent = MarketDataAgent(model_id=self.model_id)
        self.technical_analyst_agent = TechnicalAnalystAgent(model_id=self.model_id)
        self.fundamental_analyst_agent = FundamentalAnalystAgent(model_id=self.model_id)
        self.workflow: Optional[Workflow] = None

    def build_workflow(self) -> Workflow:
        """构建并返回工作流，注入已初始化的智能体实例"""
        # 输入步骤：市场数据
        def set_market_data_agent_input(step_input: StepInput) -> StepOutput:
            payload = step_input.message
            if isinstance(payload, dict):
                ticker = payload.get("ticker", "")
                start_date = payload.get("start-date") or (datetime.now() - timedelta(days=240)).strftime("%Y-%m-%d")
                end_date = payload.get("end-date") or datetime.now().strftime("%Y-%m-%d")
                news = payload.get("num-of-news") or 10
            else:
                raise ValueError("输入参数错误")

            start_date_ak = start_date.replace("-", "")
            end_date_ak = end_date.replace("-", "")
            return StepOutput(
                content=(
                    f"请收集股票 {ticker} 从 {start_date_ak} 到 {end_date_ak} 的完整市场数据，包括 {int(news)} 条新闻"
                )
            )

        market_input_step = Step(
            name="market_input_step",
            description="市场数据收集输入",
            executor=set_market_data_agent_input,
        )

        market_data_step = Step(
            name="market_data_agent",
            description="市场数据收集",
            agent=self.market_data_agent.agent,
        )

        # 输入步骤：技术分析
        def set_technical_analyst_agent_input(step_input: StepInput) -> StepOutput:
            payload = step_input.message
            if isinstance(payload, dict):
                ticker = payload.get("ticker", "")
            else:
                raise ValueError("输入参数错误")

            return StepOutput(
                content=f"对股票 {ticker} 进行技术分析"
            )

        tech_input_step = Step(
            name="tech_input_step",
            description="技术分析输入",
            executor=set_technical_analyst_agent_input,
        )

        technical_step = Step(
            name="technical_analyst_agent",
            description="技术分析",
            agent=self.technical_analyst_agent.agent,
        )

        # 输入步骤：基本面分析
        def set_fundamental_analyst_agent_input(step_input: StepInput) -> StepOutput:
            payload = step_input.message
            if isinstance(payload, dict):
                ticker = payload.get("ticker", "")
            else:
                raise ValueError("输入参数错误")
            
            return StepOutput(
                content=f"对股票 {ticker} 进行基本面分析"
            )
        
        fundamental_input_step = Step(
            name="fundamental_input_step",
            description="基本面分析输入",
            executor=set_fundamental_analyst_agent_input,
        )
        
        fundamental_step = Step(
            name="fundamental_analyst_agent",
            description="基本面分析",
            agent=self.fundamental_analyst_agent.agent,
        )

        self.workflow = Workflow(
            name="investment_analysis_workflow",
            workflow_session_state={},
            steps=[
                Steps(steps=[market_input_step, market_data_step]),
                Parallel(
                    Steps(steps=[tech_input_step, technical_step]),
                    Steps(steps=[fundamental_input_step, fundamental_step]),
                ),
            ]
            
        )
        return self.workflow

    def get_workflow_session_state(self,key:str,ticker:str):
        """获取工作流会话状态"""
        if key == "market_data":
            return self.market_data_agent.get_analysis_result(ticker)
        elif key == "technical_analysis":
            return self.technical_analyst_agent.get_analysis_result(ticker)
        elif key == "fundamental_analysis":
            return self.fundamental_analyst_agent.get_analysis_result(ticker)
        else:
            return None
        
    # 估值分析智能体和情绪分析智能体暂未接入工作流
    
def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="Agno AkShare 投资智能体系统")
    parser.add_argument("--ticker", type=str, help="股票代码")
    parser.add_argument("--start-date", type=str, help="开始日期 (YYYY-MM-DD)")
    parser.add_argument("--end-date", type=str, help="结束日期 (YYYY-MM-DD)")
    parser.add_argument("--num-of-news", type=int, help="新闻分析数量")
    parser.add_argument("--show-reasoning", action="store_true", help="显示详细推理过程")
    parser.add_argument("--model", type=str, default="deepseek-chat", help="LLM模型ID")

    args = parser.parse_args()

    # 基础目录
    os.makedirs("logs", exist_ok=True)
    os.makedirs("data", exist_ok=True)

    if not args.ticker:
        logger.error("必须提供 --ticker 参数")
        sys.exit(2)

    # 初始化系统与工作流
    system = InvestmentAnalysisSystem(model_id=args.model)
    workflow = system.build_workflow()

    # 直接传递字典作为 message（键均为短横线命名）
    message_payload = {
        "ticker": args.ticker,
        "start-date": args.start_date or None,
        "end-date": args.end_date or None,
        "num-of-news": args.num_of_news or None,
    }

    # 运行工作流（打印响应）
    workflow.print_response(
        message=message_payload,
        stream=True if args.show_reasoning else True,
        markdown=True,
    )
    
    # market_data = system.get_workflow_session_state("market_data",args.ticker)
    # _save_collected_data_to_file(args.ticker, market_data, f"market_data_agent_data_{args.ticker}_wf.json")
    technical_analysis = system.get_workflow_session_state("technical_analysis",args.ticker)
    _save_collected_data_to_file(args.ticker, technical_analysis, f"technical_analysis_agent_data_{args.ticker}_wf.json")
    fundamental_analysis = system.get_workflow_session_state("fundamental_analysis",args.ticker)
    _save_collected_data_to_file(args.ticker, fundamental_analysis, f"fundamental_analysis_agent_data_{args.ticker}_wf.json")


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        logger.info("用户中断程序")
        cleanup_on_exit()
        sys.exit(0)
    except Exception as e:
        logger.error(f"程序执行出错: {str(e)}")
        cleanup_on_exit()
        sys.exit(1)
    finally:
        # 确保程序正常退出时也执行清理
        cleanup_on_exit()